"""
Author: Morphlng
Date: 2024-11-03 22:38:50
LastEditTime: 2025-01-29 13:38:36
LastEditors: Morphlng
Description: Wrapping around Carla sensor to apply the SensorInterface.
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/carla/carla_sensor.py
"""

from __future__ import annotations

import logging
import math
import weakref
from collections import deque
from copy import deepcopy
from threading import Lock

import carla
import numpy as np

from driving_gym.misc.util import override

logger = logging.getLogger(__name__)


def is_sensor_active(sensor: carla.Sensor) -> bool:
    """Check whether the sensor is listening."""
    if not sensor:
        return False

    # Starting from 0.9.14, is_listening is a method
    tag = sensor.is_listening
    return tag() if callable(tag) else tag


def wrap_sensor(sensor: carla.Sensor) -> CarlaSensor:
    """Wrap carla sensor to SensorInterface class"""
    sensor_bp = sensor.type_id

    if sensor_bp in [
        "sensor.camera.rgb",
        "sensor.camera.semantic_segmentation",
        "sensor.camera.depth",
    ]:
        return CameraSensor(sensor)
    elif sensor_bp == "sensor.lidar.ray_cast":
        return LidarSensor(sensor)
    elif sensor_bp == "sensor.other.lane_invasion":
        return LaneInvasionSensor(sensor)
    elif sensor_bp == "sensor.other.collision":
        return CollisionSensor(sensor)
    elif sensor_bp == "sensor.other.gnss":
        return GnssSensor(sensor)
    elif sensor_bp == "sensor.other.imu":
        return IMUSensor(sensor)
    elif sensor_bp == "sensor.other.radar":
        return RadarSensor(sensor)
    elif sensor_bp == "sensor.other.obstacle":
        return ObstacleSensor(sensor)
    elif sensor_bp == "bird_eye":
        return BirdEyeSensor(sensor)
    elif sensor_bp == "sensor.other.speedometer":
        return SpeedometerSensor(sensor)

    raise ValueError(f"Unwrapped sensor type: {sensor_bp}")


class PseudoSensor:
    def __init__(self, parent: carla.Actor, sensor_type: str = "pseudo"):
        self.parent = parent
        self.type_id = sensor_type
        self.is_listening = False

    def listen(self, callback):
        self.callback = callback
        self.is_listening = True

    def stop(self):
        self.is_listening = False

    def destroy(self):
        self.is_listening = False


class CarlaSensor:
    """Basic logic for Carla sensor, should be inherited by specific sensor class"""

    def __init__(self, sensor: carla.Sensor):
        self.sensor: carla.Sensor = sensor
        self.lock = Lock()

        # data
        self.history = deque(maxlen=100)
        self.frame = -1
        self.data = None

    def reset(self):
        with self.lock:
            self.frame = -1
            self.data = None
            self.history.clear()

    def destroy(self):
        if self.sensor is None:
            return

        with self.lock:
            if is_sensor_active(self.sensor):
                self.sensor.stop()
            self.sensor.destroy()

        self.sensor = None

    def get_data(self):
        with self.lock:
            return self.data

    def get_history(self):
        with self.lock:
            return {key: value for key, value in self.history}


class CameraSensor(CarlaSensor):
    """Camera sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda image: CameraSensor._on_camera_callback(weak_self, image)
            )

    @staticmethod
    def _on_camera_callback(weak_self, image: carla.Image):
        self: CameraSensor = weak_self()
        if not self:
            return

        if self.sensor.type_id == "sensor.camera.semantic_segmentation":
            image.convert(carla.ColorConverter.CityScapesPalette)
        elif self.sensor.type_id == "sensor.camera.depth":
            image.convert(carla.ColorConverter.LogarithmicDepth)

        array = np.frombuffer(image.raw_data, dtype=np.dtype("uint8"))
        array = deepcopy(array)
        array = np.reshape(array, (image.height, image.width, 4))
        array = array[:, :, :3]
        array = array[:, :, ::-1]

        with self.lock:
            self.frame = image.frame
            self.data = array
            self.history.append((image.frame, array))


class LidarSensor(CarlaSensor):
    """Lidar sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda lidar_data: LidarSensor._on_lidar_callback(weak_self, lidar_data)
            )

    @staticmethod
    def _on_lidar_callback(weak_self, lidar_data: carla.LidarMeasurement):
        self: LidarSensor = weak_self()
        if not self:
            return

        with self.lock:
            points = np.frombuffer(lidar_data.raw_data, dtype=np.dtype("f4"))
            points = deepcopy(points)
            points = np.reshape(points, (int(points.shape[0] / 4), 4))

            self.frame = lidar_data.frame
            self.data = points
            self.history.append((lidar_data.frame, points))


class LaneInvasionSensor(CarlaSensor):
    """Lane Invasion class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        self.offlane = 0
        self.offroad = 0
        self.data = np.array([self.offlane, self.offroad])

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda event: LaneInvasionSensor._on_invasion(weak_self, event)
            )

    @override(CarlaSensor)
    def reset(self):
        with self.lock:
            self.frame = -1
            self.history.clear()
            self.offlane = 0
            self.offroad = 0
            self.data = np.array([self.offlane, self.offroad])

    @staticmethod
    def _on_invasion(weak_self, event: carla.LaneInvasionEvent):
        self: LaneInvasionSensor = weak_self()
        if not self:
            return

        with self.lock:
            self.offlane += 1
            if len(set(event.crossed_lane_markings)) == 1:
                self.offroad += 1

            self.frame = event.frame
            self.data = np.array([self.offlane, self.offroad])
            self.history.append((event.frame, self.data))


class CollisionSensor(CarlaSensor):
    """Collision sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        self.collision_vehicles = 0
        self.collision_pedestrians = 0
        self.collision_other = 0
        self.collision_id_set = set()
        self.data = np.array(
            [
                0,
                self.collision_vehicles,
                self.collision_pedestrians,
                self.collision_other,
            ]
        )

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda event: CollisionSensor._on_collision(weak_self, event)
            )

    @override(CarlaSensor)
    def reset(self):
        with self.lock:
            self.frame = -1
            self.history.clear()
            self.collision_vehicles = 0
            self.collision_pedestrians = 0
            self.collision_other = 0
            self.collision_id_set.clear()
            self.data = np.array(
                [
                    0,
                    self.collision_vehicles,
                    self.collision_pedestrians,
                    self.collision_other,
                ]
            )

    @staticmethod
    def _on_collision(weak_self, event: carla.CollisionEvent):
        self: CollisionSensor = weak_self()
        if not self:
            return

        with self.lock:
            impulse = event.normal_impulse
            intensity = math.sqrt(impulse.x**2 + impulse.y**2 + impulse.z**2)

            _cur = event.other_actor
            if _cur.id == 0 and isinstance(self.sensor.parent, carla.Walker):
                return
            else:
                if _cur.id not in self.collision_id_set:
                    self.collision_id_set.add(_cur.id)

            collided_type = type(_cur).__name__
            if collided_type == "Vehicle":
                self.collision_vehicles += 1
            elif collided_type == "Walker":
                self.collision_pedestrians += 1
            else:
                self.collision_other += 1

            self.data = np.array(
                [
                    intensity,
                    self.collision_vehicles,
                    self.collision_pedestrians,
                    self.collision_other,
                ]
            )
            self.frame = event.frame
            self.history.append((event.frame, self.data))


class GnssSensor(CarlaSensor):
    """GNSS sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        self.lat = 0.0
        self.lon = 0.0
        self.alt = 0.0

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda event: GnssSensor._on_gnss_event(weak_self, event)
            )

    @staticmethod
    def _on_gnss_event(weak_self, event: carla.GnssMeasurement):
        self: GnssSensor = weak_self()
        if not self:
            return

        with self.lock:
            self.lat = event.latitude
            self.lon = event.longitude
            self.alt = event.altitude

            self.frame = event.frame
            self.data = np.array([event.latitude, event.longitude, event.altitude])
            self.history.append((event.frame, self.data))

    @override(CarlaSensor)
    def reset(self):
        with self.lock:
            self.frame = -1
            self.data = None
            self.history.clear()
            self.lat = 0.0
            self.lon = 0.0
            self.alt = 0.0


class IMUSensor(CarlaSensor):
    """IMU sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        self.accelerometer = (0.0, 0.0, 0.0)
        self.gyroscope = (0.0, 0.0, 0.0)
        self.compass = 0.0

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda sensor_data: IMUSensor._on_imu_callback(weak_self, sensor_data)
            )

    @staticmethod
    def _on_imu_callback(weak_self, sensor_data: carla.IMUMeasurement):
        self: IMUSensor = weak_self()
        if not self:
            return

        with self.lock:
            limits = (-99.9, 99.9)
            self.accelerometer = (
                max(limits[0], min(limits[1], sensor_data.accelerometer.x)),
                max(limits[0], min(limits[1], sensor_data.accelerometer.y)),
                max(limits[0], min(limits[1], sensor_data.accelerometer.z)),
            )
            self.gyroscope = (
                max(limits[0], min(limits[1], math.degrees(sensor_data.gyroscope.x))),
                max(limits[0], min(limits[1], math.degrees(sensor_data.gyroscope.y))),
                max(limits[0], min(limits[1], math.degrees(sensor_data.gyroscope.z))),
            )
            self.compass = math.degrees(sensor_data.compass)

            self.frame = sensor_data.frame
            self.data = np.array([*self.accelerometer, *self.gyroscope, self.compass])
            self.history.append((sensor_data.frame, self.data))

    @override(CarlaSensor)
    def reset(self):
        with self.lock:
            self.frame = -1
            self.data = None
            self.history.clear()
            self.accelerometer = (0.0, 0.0, 0.0)
            self.gyroscope = (0.0, 0.0, 0.0)
            self.compass = 0.0


class RadarSensor(CarlaSensor):
    """Radar sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda radar_data: RadarSensor._on_radar_callback(weak_self, radar_data)
            )

    @staticmethod
    def _on_radar_callback(weak_self, radar_data: carla.RadarMeasurement):
        self: RadarSensor = weak_self()
        if not self:
            return

        with self.lock:
            # [depth, azimuth, altitute, velocity]
            points = np.frombuffer(radar_data.raw_data, dtype=np.dtype("f4"))
            points = deepcopy(points)
            points = np.reshape(points, (int(points.shape[0] / 4), 4))
            points = np.flip(points, 1)

            self.frame = radar_data.frame
            self.data = points
            self.history.append((radar_data.frame, points))


class ObstacleSensor(CarlaSensor):
    """Obstacle sensor class wrapper"""

    def __init__(self, sensor: carla.Sensor):
        super().__init__(sensor)

        self.data = np.array([-1, 0.0])

        if not is_sensor_active(self.sensor):
            weak_self = weakref.ref(self)
            self.sensor.listen(
                lambda event: ObstacleSensor._on_obstacle_callback(weak_self, event)
            )

    @staticmethod
    def _on_obstacle_callback(weak_self, event: carla.ObstacleDetectionEvent):
        self: ObstacleSensor = weak_self()
        if not self:
            return

        with self.lock:
            distance = event.distance
            obstacle_id = event.other_actor.id
            self.frame = event.frame
            self.data = np.array([obstacle_id, distance])
            self.history.append((event.frame, self.data))


class BirdEyeSensor(CarlaSensor):
    """Bird-eye view sensor class wrapper"""

    def __init__(self, sensor):
        super().__init__(sensor)

        self.parent: carla.Actor = sensor.parent
        self.world: carla.World = self.sensor._world

        # Create first frame to avoid empty data
        self.data = self._generate_data()

    @override(CarlaSensor)
    def reset(self):
        self.frame = -1
        self.data = self._generate_data()
        self.history.clear()

    @override(CarlaSensor)
    def get_data(self):
        self.data = self._generate_data()
        self.frame = self.world.get_snapshot().frame
        self.history.append((self.frame, self.data))
        return self.data

    @override(CarlaSensor)
    def destroy(self):
        self.frame = -1
        self.data = None
        self.history.clear()
        self.parent = None
        self.world = None

    def _generate_data(self):
        raw_data = self.sensor.produce(self.parent)
        rgb_data = self.sensor.as_rgb(raw_data)
        return rgb_data


class SpeedometerSensor(CarlaSensor):
    """Sensor to measure the speed of the vehicle."""

    @override(CarlaSensor)
    def reset(self):
        self.frame = -1
        self.data = 0.0
        self.history.clear()

    @override(CarlaSensor)
    def get_data(self):
        velocity = self.sensor.parent.get_velocity()
        transform = self.sensor.parent.get_transform()
        return self._get_forward_speed(transform, velocity)

    def _get_forward_speed(self, transform: carla.Transform, velocity: carla.Vector3D):
        """Convert the vehicle transform directly to forward speed"""
        vel_np = np.array([velocity.x, velocity.y, velocity.z])
        pitch = np.deg2rad(transform.rotation.pitch)
        yaw = np.deg2rad(transform.rotation.yaw)
        orientation = np.array(
            [np.cos(pitch) * np.cos(yaw), np.cos(pitch) * np.sin(yaw), np.sin(pitch)]
        )
        speed = np.dot(vel_np, orientation)
        return np.array([speed])
